# -*- coding: utf-8 -*-
import sqlite3
import json
import re
import os
from typing import Optional, List, Dict, Any, Union
from ..core.logger import get_logger

logger = get_logger(__name__)

class FinanceDatabaseTool:
    def __init__(self):
        # 设置数据库路径
        self.db_path = os.path.join(
            os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
            "dataset", "sqlite", "博金杯比赛数据.db"
        )
        logger.info(f"数据库路径: {self.db_path}")

    def execute_query(self, query: str) -> List[Dict[str, Any]]:
        """执行SQL查询并返回结果"""
        try:
            # 处理查询中的带括号列名
            sanitized_query = self._sanitize_query(query)
            logger.info(f"执行查询: {sanitized_query}")
            
            # 连接数据库
            conn = sqlite3.connect(self.db_path)
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            
            # 执行查询
            cursor.execute(sanitized_query)
            results = cursor.fetchall()
            
            # 转换为字典列表
            dict_results = [dict(row) for row in results]
            
            # 关闭连接
            conn.close()
            
            logger.info(f"查询结果数量: {len(dict_results)}")
            return dict_results
        except Exception as e:
            logger.error(f"查询执行错误: {e}")
            return []

    def _sanitize_query(self, query: str) -> str:
        """清理和修复SQL查询，特别是处理带括号的列名"""
        # 使用简单直接的方式处理带括号的列名
        # 针对港股日行情表中常见的带括号列名进行特殊处理
        column_names = [
            "最高价(元)", "最低价(元)", "收盘价(元)", "昨收盘(元)", "今开盘(元)",
            "成交量(股)", "成交金额(元)"
        ]
        
        # 为每个带括号的列名添加双引号
        for column in column_names:
            if column in query:
                query = query.replace(column, f'"{column}"')
        
        # 同时确保表名也被正确引用（如果使用了单引号）
        query = query.replace("'港股票日行情表'", '"港股票日行情表"')
        
        return query

    def get_table_schema(self, table_name: str = None) -> Union[Dict[str, List[Dict[str, str]]], List[Dict[str, str]]]:
        """
        获取表的结构
        
        Args:
            table_name: 可选，如果提供，则返回特定表的结构；否则返回所有表的结构
        
        Returns:
            如果指定了表名，返回该表的列信息列表；否则返回所有表的结构字典
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # 如果指定了表名，只获取该表的结构
            if table_name:
                cursor.execute(f"PRAGMA table_info({table_name});")
                columns = cursor.fetchall()
                
                result = []
                for column in columns:
                    result.append({
                        "name": column[1],
                        "type": column[2]
                    })
            else:
                # 否则获取所有表的结构
                cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
                tables = [table[0] for table in cursor.fetchall()]
                
                schema = {}
                for table in tables:
                    cursor.execute(f"PRAGMA table_info({table});")
                    columns = cursor.fetchall()
                    
                    schema[table] = []
                    for column in columns:
                        schema[table].append({
                            "name": column[1],
                            "type": column[2]
                        })
                
                result = schema
            
            conn.close()
            return result
        except Exception as e:
            logger.error(f"获取表结构错误: {e}")
            return [] if table_name else {}

    def get_available_tables(self) -> List[str]:
        """获取所有可用的表名"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
            tables = [table[0] for table in cursor.fetchall()]
            
            conn.close()
            return tables
        except Exception as e:
            logger.error(f"获取表名错误: {e}")
            return []
    
    def get_table_names(self) -> List[str]:
        """获取数据库中所有表名（为了保持兼容性添加的方法）"""
        # 直接调用现有的get_available_tables方法
        return self.get_available_tables()
        
    def get_table_sample_data(self, table_name: str, limit: int = 5) -> List[Dict[str, Any]]:
        """
        获取指定表的示例数据
        
        Args:
            table_name: 表名
            limit: 返回的记录数，默认为5
        
        Returns:
            表的示例数据列表
        """
        try:
            conn = sqlite3.connect(self.db_path)
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            
            # 获取表的示例数据
            query = f"SELECT * FROM {table_name} LIMIT ?"
            cursor.execute(query, (limit,))
            results = cursor.fetchall()
            
            # 转换为字典列表
            sample_data = [dict(row) for row in results]
            
            conn.close()
            return sample_data
        except Exception as e:
            logger.error(f"获取表示例数据错误: {e}")
            return []

# 创建工具实例
finance_db_tool = FinanceDatabaseTool()

# 定义供Agent调用的函数
def execute_finance_query(query: str) -> str:
    """
    执行金融数据库查询
    
    Args:
        query: SQL查询语句
    
    Returns:
        str: 查询结果的JSON字符串
    """
    try:
        results = finance_db_tool.execute_query(query)
        # 限制返回的结果数量，避免数据过大
        if len(results) > 100:
            results = results[:100]
            results.append({"message": "结果已限制为前100条记录"})
        return json.dumps(results, ensure_ascii=False, default=str)
    except Exception as e:
        logger.error(f"执行金融查询时出错: {e}")
        return json.dumps({"error": str(e)}, ensure_ascii=False)